#include <atomic>
#include <chrono>
#include <iostream>
#include <thread>
#include <vector>

#include "data_dumper.h"

constexpr int kNumThreads = 10;
constexpr int kNumPushesPerThread = 10000;
constexpr int kSleepTimeUs = 10;

std::atomic<int> total_pushed{0};

void PushData(DataDumper& dumper) {
  for (int i = 0; i < kNumPushesPerThread; ++i) {
    std::string op_name =
        "Default/"
        "TestXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX"
        "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX/"
        "MatMul_Op_" +
        std::to_string(i);
    auto data = std::make_unique<MindSporeOpRangeData>(0,     // device_id
                                                       0,     // tag
                                                       1,     // thread_id
                                                       2,     // flow_id
                                                       3,     // step
                                                       4,     // start_time_ns
                                                       5,     // end_time_ns
                                                       6,     // process_id
                                                       7,     // module_index
                                                       8,     // event_index
                                                       9,     // stage_index
                                                       10,    // level
                                                       true,  // is_async
                                                       true,  // is_graph_data
                                                       true,  // is_stage
                                                       op_name);
    dumper.Report(std::move(data));
    total_pushed.fetch_add(1, std::memory_order_relaxed);
    std::this_thread::sleep_for(std::chrono::microseconds(kSleepTimeUs));
  }
}

int main() {
  auto start_time = std::chrono::high_resolution_clock::now();
  DataDumper& dumper = DataDumper::GetInstance();
  dumper.Init("./tmp", 100 * 1024);
  dumper.Start();

  std::vector<std::thread> threads;
  for (int i = 0; i < kNumThreads; ++i) {
    threads.emplace_back(PushData, std::ref(dumper));
  }

  for (auto& t : threads) {
    t.join();
  }

  dumper.Stop();

  // 检查是否有数据丢失
  int expected_total = kNumThreads * kNumPushesPerThread;
  int actual_total = total_pushed.load(std::memory_order_relaxed);
  std::cout << "thread num: " << kNumThreads
            << ", push num per thread: " << kNumPushesPerThread << std::endl;
  std::cout << "expected total: " << expected_total << std::endl;
  std::cout << "actual total: " << actual_total << std::endl;

  auto end_time = std::chrono::high_resolution_clock::now();
  std::cout << "cost time: "
            << std::chrono::duration_cast<std::chrono::milliseconds>(end_time -
                                                                     start_time)
                   .count()
            << "ms" << std::endl;

  return 0;
}